{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ajuste Fino Supervisionado com SFTTrainer\n",
    "\n",
    "Este caderno demonstra como ajustar o modelo `HuggingFaceTB/SmolLM2-135M` usando o `SFTTrainer` do móduli `trl`. As células do caderno são executadas e vão ajustar finamente o modelo. Você pode selecionar sua dificuldade experimentando diferentes conjuntos de dados.\n",
    "\n",
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Exercício: Ajuste fino do SmolLM2 com SFTTrainer</h2>\n",
    "    <p>Pegue um conjunto de dados do hub do Hugging Face e faça um ajuste fino em um modelo. </p> \n",
    "    <p><b>Níveis de Dificuldade</b></p>\n",
    "    <p>🐢 Use o conjunto de dados `HuggingFaceTB/smoltalk`.</p>\n",
    "    <p>🐕 Experimente o conjunto de dados `bigcode/the-stack-smol` e ajuste um modelo de geração de código no subconjunto específico `data/python`.</p>\n",
    "    <p>🦁 Selecione um conjunto de dados relacionado a um caso de uso do mundo real que seja do seu interesse.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instale os requisitos no Google Colab\n",
    "# !pip install transformers datasets trl huggingface_hub\n",
    "\n",
    "# Autentique ao Hugging Face\n",
    "\n",
    "from huggingface_hub import login\n",
    "login()\n",
    "\n",
    "# Por conveniência, você pode criar uma variável de ambiente que contém seu token do hub como HF_TOKEN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importe módulos necessários\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "from trl import SFTConfig, SFTTrainer, setup_chat_format\n",
    "import torch\n",
    "\n",
    "device = (\n",
    "    \"cuda\"\n",
    "    if torch.cuda.is_available()\n",
    "    else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    ")\n",
    "\n",
    "# Carregue o modelo e o tokenizador\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=model_name\n",
    ").to(device)\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)\n",
    "\n",
    "# Configure o formato de bate-papo\n",
    "model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)\n",
    "\n",
    "# Defina nosso nome para o ajuste fino a ser salvo e/ou carregado em um local\n",
    "finetune_name = \"SmolLM2-FT-MyDataset\"\n",
    "finetune_tags = [\"smol-course\", \"module_1\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gerar com o modelo base\n",
    "\n",
    "Aqui vamos experimentar o modelo base que não tem um modelo de bate-papo. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vamos testar o modelo base antes do treinamento\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Formato com o modelo de bate-papo\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Gerando a resposta\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "outputs = model.generate(**inputs, max_new_tokens=100)\n",
    "print(\"Before training:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparação do Conjunto de Dados\n",
    "\n",
    "Carregaremos um conjunto de dados de amostra e o formataremos para treinamento. O conjunto de dados deve ser estruturado com pares de input-output, em que cada input é um prompt e o output é a resposta esperada do modelo.\n",
    "\n",
    "**TRL formatará as mensagens de entrada com base nos modelos de bate-papo.** Elas precisam ser representadas como uma lista de dicionários com as chaves: `role` e `content`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Carregue um conjunto de dados de amostra\n",
    "from datasets import load_dataset\n",
    "\n",
    "# TODO: Defina seu próprio conjunto de dados e configurações usando os parâmetros do caminho e do nome\n",
    "ds = load_dataset(path=\"HuggingFaceTB/smoltalk\", name=\"everyday-conversations\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: 🦁 Se o seu conjunto de dados não estiver em um formato que o TRL possa converter para o modelo de chat, você precisará processá-lo. Consulte o [módulo](../chat_templates.md)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuração do SFTTrainer\n",
    "\n",
    "O `SFTTrainer` é configurado com vários parâmetros que controlam o processo de treinamento. Eles incluem o número de etapas de treinamento, o tamanho do lote, a taxa de aprendizado e a estratégia de avaliação. Ajuste esses parâmetros com base em seus requisitos específicos e recursos computacionais."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure o SFTTrainer\n",
    "sft_config = SFTConfig(\n",
    "    output_dir=\"./sft_output\",\n",
    "    max_steps=1000,  # Adjust based on dataset size and desired training duration\n",
    "    per_device_train_batch_size=4,  # Set according to your GPU memory capacity\n",
    "    learning_rate=5e-5,  # Common starting point for fine-tuning\n",
    "    logging_steps=10,  # Frequency of logging training metrics\n",
    "    save_steps=100,  # Frequency of saving model checkpoints\n",
    "    evaluation_strategy=\"steps\",  # Evaluate the model at regular intervals\n",
    "    eval_steps=50,  # Frequency of evaluation\n",
    "    use_mps_device=(\n",
    "        True if device == \"mps\" else False\n",
    "    ),  # Use MPS for mixed precision training\n",
    "    hub_model_id=finetune_name,  # Set a unique name for your model\n",
    ")\n",
    "\n",
    "# Inicie o SFTTrainer\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=sft_config,\n",
    "    train_dataset=ds[\"train\"],\n",
    "    eval_dataset=ds[\"test\"],\n",
    ")\n",
    "\n",
    "# TODO: 🦁 🐕 alinhe os parâmetros do SFTTrainer com o conjunto de dados escolhido. Por exemplo, se estiver usando o conjunto de dados `bigcode/the-stack-smol`, será necessário escolher a coluna `content`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Treinando o modelo\n",
    "\n",
    "Com o trainer configurado, podemos prosseguir com o treinamento do modelo. O processo de treinamento envolverá a iteração do conjunto de dados, o cálculo da perda e a atualização dos parâmetros do modelo para minimizar essa perda."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Treine o modelo\n",
    "trainer.train()\n",
    "\n",
    "# Salve o model\n",
    "trainer.save_model(f\"./{finetune_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.push_to_hub(tags=finetune_tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style='background-color: lightblue; padding: 10px; border-radius: 5px; margin-bottom: 20px; color:black'>\n",
    "    <h2 style='margin: 0;color:blue'>Exercício bônus: Gerar com o modelo ajustado</h2>\n",
    "    <p>🐕 Use o ajuste fino para modelar e gerar uma resposta, assim como no exemplo básico.</p>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Teste o modelo ajustado no mesmo prompt\n",
    "\n",
    "# Vamos testar o modelo básico antes do treinamento\n",
    "prompt = \"Write a haiku about programming\"\n",
    "\n",
    "# Formatar com o modelo\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "\n",
    "# Gerar resposta\n",
    "inputs = tokenizer(formatted_prompt, return_tensors=\"pt\").to(device)\n",
    "\n",
    "# TODO: use o ajuste fino para modelar e gerar uma resposta, assim como no exemplo básico."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 💐 Você conseguiu!\n",
    "\n",
    "Este caderno fornece um guia passo-a-passo para o ajuste fino do modelo `HuggingFaceTB/SmolLM2-135M` usando o `SFTTrainer`. Seguindo essas etapas, você pode adaptar o modelo para executar tarefas específicas com mais eficiência. Se quiser continuar trabalhando neste curso, aqui estão as etapas que você pode experimentar:\n",
    "\n",
    "- Experimente este caderno com uma dificuldade maior\n",
    "- Revisar o PR de um colega\n",
    "- Melhorar o material do curso por meio de uma Issue ou PR."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
